/**
 * @file facelib.cpp
 * @author lynxi
 * @brief
 * @version 1.0
 * @date 2024-04-30
 *
 * Copyright:
 * © 2024 北京灵汐科技有限公司 版权所有。
 * 注意：以下内容均为北京灵汐科技有限公司原创，未经本公司允许，不得转载，否则将视为侵权；对于不遵守此声明或者其他违法使用以下内容者，本公司依法保留追究权。
 * © 2024 Lynxi Technologies Co., Ltd. All rights reserved.
 * NOTICE: All information contained here is, and remains the property of Lynxi.
 * This file can not be copied or distributed without the permission of Lynxi
 * Technologies Co., Ltd.
 *
 */

#include <lyn_api.h>
#include <nlohmann/json.hpp>
#include <string>
#include "imageDecoder.h"
#include "ipeParamFaceRecognition.h"
#include "CLI/CLI11.hpp"

using namespace std;

int main(int argc, char *argv[])
{
  CLI::App app;
  string faceImagePath;
  string configFile;
  string modelPath;
  int deviceId;

  app.add_option("-i,--face_image_path", faceImagePath, "input face image path")
      ->default_val("../data/face_recognition_with_image_sample/face_images");
  app.add_option("-c,--config_file", configFile, "input face config file")
      ->default_val("../data/face_recognition_with_image_sample/face_config.json");
  app.add_option("-m,--model_path", modelPath, "face feature model path")
      ->default_val("../model/w600k_mbf_b1_h112_w112_c3_iuint8_ofloat16_0000/Net_0");
  app.add_option("-d,--device_id", deviceId, "lynxi device id")
      ->default_val(0);

  CLI11_PARSE(app, argc, argv);

  if (!IsFileExist(faceImagePath))
  {
    cout << "error : input face image path is not exsit." << endl;
    return 0;
  }

  if (!IsFileExist(configFile))
  {
    cout << "error : input face config file is not exsit." << endl;
    return 0;
  }

  if (!IsFileExist(modelPath))
  {
    cout << "error : face feature model path is not exsit." << endl;
    return 0;
  }

  if (!IsValidDevice(deviceId))
  {
    cout << "error : device id is invalid." << endl;
    return 0;
  }

  ifstream fJson(configFile);
  stringstream ss;
  ss << fJson.rdbuf();
  nlohmann::json js = nlohmann::json::parse(ss.str());

  string featurePath = js["features"].get<string>();
  vector<string> images = js["images"].get<vector<string>>();

  ofstream featuresBin(featurePath);

  // 2. 创建 context 与 stream
  lynContext_t ctx = nullptr;
  CHECK_ERR(lynCreateContext(&ctx, deviceId));
  CHECK_ERR(lynRegisterErrorHandler(StreamErrorHandler, nullptr));
  lynStream_t stream = nullptr;
  CHECK_ERR(lynCreateStream(&stream));

  // 3. 初始化
  ModelInfo modelInfo;
  modelInfo.LoadModelByPath(modelPath.c_str());

  IpeParamW600k_mbf ipe(modelInfo.width, modelInfo.height);

  // 图片解码缓存
  uint8_t *pDecodeBufOut = nullptr;
  CHECK_ERR(lynMalloc((void **)&pDecodeBufOut, 1920 * 1080 * 3 / 2));
  // ipe图像缓存
  uint8_t *pIpeBufOut = nullptr;
  CHECK_ERR(lynMalloc((void **)&pIpeBufOut, modelInfo.inputSize));
  // dev推理结果缓存
  uint8_t *pDevInferBufOut = nullptr;
  CHECK_ERR(lynMalloc((void **)&pDevInferBufOut, modelInfo.outputSize));

  // 循环将人脸图片录入人脸库中
  cout << "store in face lib ..." << endl;
  for (auto &image : images)
  {
    // 4. 初始化解码图片类
    ImageDecoder imageDec;
    imageDec.Init(faceImagePath + "/" + image, false);
    lynImageInfo_t imgInfo;
    imageDec.GetImageDecInfo(&imgInfo);

    // 5. 同步解码图片到 Device 侧
    lynCodecBuf_t decodeImg;
    decodeImg.size = imgInfo.output.predictBufSize;
    decodeImg.data = pDecodeBufOut;
    imageDec.DecodeImageToDevice(stream, decodeImg.data);

    // 6. IPE 处理
    ipe.SetImgInfo(imgInfo.output.width, imgInfo.output.height,
                   LYN_PIX_FMT_NV12);
    ipe.CalcParam(stream, decodeImg.data, pIpeBufOut);

    // 7. 调用 APU 推理接口
    CHECK_ERR(lynExecuteModelAsync(stream, modelInfo.model, pIpeBufOut,
                                   pDevInferBufOut, modelInfo.batchSize));

    // host推理结果缓存
    uint8_t *pHostInferBufOut = new uint8_t[modelInfo.outputSize];
    CHECK_ERR(lynMemcpyAsync(stream, pHostInferBufOut, pDevInferBufOut,
                             modelInfo.outputSize, ServerToClient));
    SetCallback(stream, [pHostInferBufOut, modelInfo, &featuresBin]
                {
      featuresBin.write((const char *)pHostInferBufOut, modelInfo.outputSize);
      delete[] pHostInferBufOut; });
    CHECK_ERR(lynSynchronizeStream(stream));
  }

  cout << "store over" << endl;

  // 8. 销毁资源
  CHECK_ERR(lynDestroyStream(stream));
  CHECK_ERR(lynFree(pDevInferBufOut));
  CHECK_ERR(lynFree(pIpeBufOut));
  CHECK_ERR(lynFree(pDecodeBufOut));
  modelInfo.UnLoadModel();
  CHECK_ERR(lynDestroyContext(ctx));

  return 0;
}
